from gc import callbacks
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import optimizers
from ResNeXt import resnext  # 导入模型
import os
import matplotlib.pyplot as plt

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'  # 调用GPU训练
os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # 使用当前设备的第一块GPU

gpus = tf.config.experimental.list_physical_devices("GPU")

# 设置GPU内存占用，根据网络模型大小占用相应的内存
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)
        exit(-1)

# ------------------------------------- #
# 加载数据集，batchsz太大会导致显存爆炸
# ------------------------------------- #
filepath = 'D:/deeplearning/test/数据集/交通标志/new_data/'  # 数据集所在文件夹
height = 224  # 输入图象的高
width = 224  # 输入图象的宽
batchsz = 8  # 每个batch处理32张图片
checkData = False  # 查看数据集划分的信息
plotShow = True  # 绘制图像
checkDataAgain = False  # 预处理之后是否再次查看数据集信息
# ------------------------------------- #
# 网络模型结构
# ------------------------------------- #
input_shape = (224,224,3)  # 输入图象的shape
classes = 4  # 图像分类的类别数
checkNet = False  # 是否查看网络架构
# ------------------------------------- #
# 网络训练
# ------------------------------------- #
weights_dir = 'save_weights'  # 权重文件保存的文件夹路径
learning_rate = 0.0005  # adam优化器的学习率
epochs = 10  # 训练迭代次数


#（1）加载数据集
def get_data(height, width, batchsz):
    
    # 训练集数据
    train_ds = keras.preprocessing.image_dataset_from_directory(
        directory = filepath + 'train',  # 训练集图片所在文件夹
        label_mode = 'categorical',  # onehot编码
        image_size = (height, width),  # 输入图象的size
        batch_size = batchsz)  # 每批次训练32张图片

    # 验证集数据
    val_ds = keras.preprocessing.image_dataset_from_directory(
        directory = filepath + 'val', 
        label_mode = 'categorical', 
        image_size = (height, width),  
        batch_size = batchsz)  


    # 返回划分好的数据集
    return train_ds, val_ds

# 读取数据集
train_ds, val_ds = get_data(height, width, batchsz) 

#（2）查看数据集信息
def check_data(train_ds):  # 传入训练集数据集
    
    # 查看数据集有几个分类类别
    class_names = train_ds.class_names
    print('classNames:', class_names)

    # 查看数据集的shape, x代表图片数据, y代表分类类别数据
    sample = next(iter(train_ds))  # 生成迭代器,每次取出一个batch的数据
    print('x_batch.shape:', sample[0].shape, 'y_batch.shape:', sample[1].shape)
    print('前五个目标值:', sample[1][:5])

# 是否查看数据集信息
if checkData is True:
    check_data(train_ds)


#（3）查看图像
def plot_show(train_ds):

    # 生成迭代器,每次取出一个batch的数据
    sample = next(iter(train_ds))  # sample[0]图像信息, sample[1]标签信息
    # 显示前5张图
    for i in range(5):
        plt.subplot(1,5,i+1)  # 在一块画板的子画板上绘制1行5列
        plt.imshow(sample[0][i])  # 图像的像素值压缩到0-1
        plt.xticks([])  # 不显示xy坐标刻度
        plt.yticks([])
    plt.show()

# 是否展示图像信息
if plotShow is True:
    plot_show(train_ds)


#（4）数据预处理
def processing(x,y):  # 定义预处理函数
    x = tf.cast(x, dtype=tf.float32) / 255.0  # 图片转换为tensor类型，并归一化
    y = tf.cast(y, dtype=tf.int32)  # 分类标签转换成tensor类型
    return x,y

# 对所有数据预处理
train_ds = train_ds.map(processing).shuffle(10000)  # map调用自定义预处理函数, shuffle打乱数据集
val_ds = val_ds.map(processing)


#（5）查看预处理后的数据是否正确
def check_data_again(train_ds):  # 传入训练集数据集

    sample = next(iter(train_ds))  # 生成迭代器,每次取出一个batch的数据
    print('-------after preprocessing-------')
    print('x_batch.shape:', sample[0].shape, 'y_batch.shape:', sample[1].shape)
    print('前五个目标值:', sample[1][:5])

# 是否查看数据集信息
if checkDataAgain is True:
    check_data_again(train_ds)


#（6）导入网络模型
model = resnext(input_shape=input_shape, # 网络的输入图像的size
                classes=classes)  # 分类数

# 查看网络构架
if checkNet is True:
    model.summary()


#（7）保存权重文件
if not os.path.exists(weights_dir):  # 判断当前文件夹下有没有一个叫save_weights的文件夹
    os.makedirs(weights_dir)  # 如果没有就创建一个


#（8）模型编译
opt = optimizers.Adam(learning_rate=learning_rate)  # 设置Adam优化器

model.compile(optimizer=opt, #学习率
              loss=keras.losses.CategoricalCrossentropy(from_logits=True), # 交叉熵损失，logits层先经过softmax
              metrics=['accuracy']) #评价指标


#（9）定义回调函数，一个列表
# 保存模型参数
callbacks = [keras.callbacks.ModelCheckpoint(filepath = 'save_weights/resnext.h5',  # 参数保存的位置
                                            save_best_only = True,  # 保存最佳参数
                                            save_weights_only = True,  # 只保存权重文件
                                            monitor = 'val_loss')]  # 通过验证集损失判断是否是最佳参数


#（10）模型训练，history保存训练信息
history = model.fit(x = train_ds,  # 训练集
                    validation_data = val_ds,  # 验证集
                    epochs = epochs,  #迭代30次
                    callbacks = callbacks) 


#（11）获取训练信息
history_dict = history.history  # 获取训练的数据字典
train_loss = history_dict['loss']  # 训练集损失
train_accuracy = history_dict['accuracy']  # 训练集准确率
val_loss = history_dict['val_loss']  # 验证集损失
val_accuracy = history_dict['val_accuracy']  # 验证集准确率

#（12）绘制训练损失和验证损失
plt.figure()
plt.plot(range(epochs), train_loss, label='train_loss')  # 训练集损失
plt.plot(range(epochs), val_loss, label='val_loss')  # 验证集损失
plt.legend()  # 显示标签
plt.xlabel('epochs')
plt.ylabel('loss')

#（13）绘制训练集和验证集准确率
plt.figure()
plt.plot(range(epochs), train_accuracy, label='train_accuracy')  # 训练集准确率
plt.plot(range(epochs), val_accuracy, label='val_accuracy')  # 验证集准确率
plt.legend()
plt.xlabel('epochs')
plt.ylabel('accuracy')

